"""
@Time: 2020/12/29 下午 6:33
@Author: jinzhuan
@File: ace2005.py
@Desc: 
"""
import os
from ..loader import Loader
from cognlp.utils.util import load_json


class ACE2005RelationLoader(Loader):

    def __init__(self):
        super().__init__()

    def _load(self, path):
        data = load_json(path)
        dataset = []
        for item in data:
            sentence = item['sentence']
            words = item['tokens']
            entity_mentions = item['golden-entity-mentions']
            if len(entity_mentions) == 0:
                continue
            relation_mentions = item['golden-relation-mentions']
            for relation_mention in relation_mentions:
                self.label_set.add(relation_mention['relation-type'])
            dataset.append({'sentence': sentence, 'words': words, 'entity_mentions': entity_mentions,
                            'relation_mentions': relation_mentions})
        return dataset

    def load_all(self, path):
        train_path = os.path.join(path, 'en-train.json')
        dev_path = os.path.join(path, 'en-dev.json')
        test_path = os.path.join(path, 'en-test.json')
        return self._load(train_path), self._load(dev_path), self._load(test_path)
